#include <stdio.h>
#include <iostream>
#include <string>
#include <fstream>
#include <time.h>
#include <sys/time.h>
#include <opencv2/opencv.hpp>
#include <opencv2/nonfree/nonfree.hpp>
#include "opencv2/nonfree/features2d.hpp"
#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/calib3d/calib3d.hpp"
#include <omp.h>
using namespace std;
using namespace cv;

double timeDiff(struct timeval t1, struct timeval t2);
void FindMatching(vector<vector<DMatch> > matches, vector<KeyPoint> keypoints1,
		vector<KeyPoint> keypoints2, vector<Point2f>* srcPoints,
		vector<Point2f>* destPoints);
Mat CalculateHomography(int n, vector<Point2f>* p1, vector<Point2f>* p2);
Mat PerformRANSAC(int ID, double decisionThreshold, int numberOfRandomSeeds,
		int numberOfTrials, float acceptThresholdRatio,
		const vector<Point2f>* srcPoints, const vector<Point2f>* destPoints,
		vector<Point2f>* inlinerSrcPoints, vector<Point2f>* inlinerDestPoints);
Mat CreateMatchingImage(Mat image1, Mat image2,
		const vector<Point2f>* srcPoints, const vector<Point2f>* destPoints,
		const vector<Point2f>* inlinerSrcPoints,
		const vector<Point2f>* inlinerDestPoints);

int main(int argc, char** argv) {
	struct timeval StartTime;
	struct timeval EndTime;
	struct timeval Time1;
	struct timeval Time2;
	double duration;
	int i;


	if (argc != 3) {
		printf("Usage:./t1_ImageMosacing <9/27> <Number of Threads>\n");
		return EXIT_FAILURE;
	}
        if (atoi(argv[1]) != 9 && atoi(argv[1]) != 27) {
                printf("Usage: ./t1_ImageMosacing <9/27>\n");
                return EXIT_FAILURE;
        }
	int NUM_IMAGES = atoi(argv[1]);

	Mat images[NUM_IMAGES];
	Mat matchImages[NUM_IMAGES - 1];
	Mat correctedImages[NUM_IMAGES];
	vector<KeyPoint> interestPoints[NUM_IMAGES];
	Mat descriptors[NUM_IMAGES];
	vector<vector<DMatch> > matches[NUM_IMAGES - 1];
	vector<Point2f> srcPoints[NUM_IMAGES - 1];
	vector<Point2f> destPoints[NUM_IMAGES - 1];
	vector<Point2f> inlinerSrcPoints[NUM_IMAGES - 1];
	vector<Point2f> inlinerDestPoints[NUM_IMAGES - 1];
	Mat H[NUM_IMAGES - 1];
	Mat corrH[NUM_IMAGES];

	Mat imageMasks[NUM_IMAGES];
	const int IMAGE_WIDTH_RATIO = 5;
	const int IMAGE_HEIGHT_RATIO = 3;
	const int OFF_WIDTH = 2;
	const int OFF_HEIGHT = 1;

	Mat blankMask;
	Mat finalImage;




	omp_set_num_threads(atoi(argv[2]));
	printf("Number of threads is set to %d\n", atoi(argv[2]));

	cv::SurfFeatureDetector detector(300);
	cv::SurfDescriptorExtractor extractor;
	cv::DescriptorMatcher* matcher = new BFMatcher(NORM_L2, true);

        string myString1 = "Pictures/9ImageSet/";
        if (NUM_IMAGES == 27) {
                myString1 = "Pictures/27ImageSet/";
        }


	string myString2 = ".JPG";
	string myString3;

	gettimeofday(& StartTime, NULL);
	//Read images
	gettimeofday(& Time1, NULL);
	for (i = 0; i < NUM_IMAGES; i++) {
		std::stringstream out;
		out << (i + 1);
		myString3 = myString1 + out.str() + myString2;
		images[i] = imread(myString3, CV_LOAD_IMAGE_COLOR);

		if (!images[i].data) {
			printf("Error reading image %d\n", i + 1);
			return EXIT_FAILURE;
		}
	}
	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2);
	printf("Time to read %d images takes ,%.0f, ms\n", NUM_IMAGES, duration);

	//Detect interest points of the images
	gettimeofday(& Time1, NULL);
#pragma omp parallel for
	for (i = 0; i < NUM_IMAGES; i++) {
		detector.detect(images[i], interestPoints[i]);
	}
	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2); 
	printf("Time to extract interest points from %d images takes ,%.0f, ms\n", NUM_IMAGES, duration);

	//Compute Descriptor for interest Points
	gettimeofday(& Time1, NULL);
#pragma omp parallel for
	for (i = 0; i < NUM_IMAGES; i++) {
		extractor.compute(images[i], interestPoints[i], descriptors[i]);
	}
	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2); 
	printf("Time to compute descriptors for %d images takes ,%.0f, ms\n", NUM_IMAGES, duration);

	//Find Best Correspondences
	gettimeofday(& Time1, NULL);
#pragma omp parallel for
	for (i = 0; i < NUM_IMAGES - 1; i++) {
		matcher->knnMatch(descriptors[i], descriptors[i + 1], matches[i], 1);
	}
	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2); 
	printf("Time to compute correspondences for %d image pairs takes ,%.0f, ms\n", NUM_IMAGES - 1, duration);

	//Store correspondence in srcPoints and destPoints vectors
	gettimeofday(& Time1, NULL);
#pragma omp parallel for
	for (i = 0; i < NUM_IMAGES - 1; i++) {
		FindMatching(matches[i], interestPoints[i], interestPoints[i + 1],
				&srcPoints[i], &destPoints[i]);
	}
	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2); 
	printf("Time to store correspondences for %d image pairs takes ,%.0f, ms\n", NUM_IMAGES - 1, duration);

	//Run RANSAC
	gettimeofday(& Time1, NULL);
	srand(time(NULL));
	for (i = 0; i < NUM_IMAGES - 1; i++) {
		H[i] = PerformRANSAC(i, 3, 5, 20000, 0.1, &srcPoints[i], &destPoints[i], &inlinerSrcPoints[i], &inlinerDestPoints[i]);
	}
	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2); 
	printf("Time to run RANSAC for %d image pairs takes ,%.0f, ms\n", NUM_IMAGES - 1, duration);

	//Create image pairs with inliner and outliner matching
	gettimeofday(& Time1, NULL);
#pragma omp parallel for
	for (i = 0; i < NUM_IMAGES - 1; i++) {
		matchImages[i] = CreateMatchingImage(images[i], images[i + 1], &srcPoints[i], &destPoints[i], &inlinerSrcPoints[i], &inlinerDestPoints[i]);
	}
	gettimeofday(& Time1, NULL);
	duration = timeDiff(Time1, Time2);
	printf("Time to create %d matching image pairs takes ,%.0f, ms\n", NUM_IMAGES - 1, duration);

	//Calculate homographys with respect to the center image
	gettimeofday( &Time1, NULL);
	Mat offset = (Mat_<double>(3, 3) << 1, 0, -OFF_WIDTH * images[0].cols, 0, 1, -OFF_HEIGHT *images[0].rows, 0, 0, 1);

	/*
	 * H
	 * 0 - 1_2
	 * 1 - 2_3
	 * 2 - 3_4
	 * 3 - 4_5
	 * 4 - 5_6
	 * 5 - 6_7
	 * 6 - 7_8
	 * 7 - 8_9
	 */

	/*
	 * corrH
	 * 0 - 1_5
	 * 1 - 2_5
	 * 2 - 3_5
	 * 3 - 4_5
	 * 4 - 5-5
	 * 5 - 5_6
	 * 6 - 5_7
	 * 7 - 5_8
	 * 8 - 5_9
	 */

	corrH[NUM_IMAGES/2 - 1] = H[NUM_IMAGES/2 - 1];
	for(i = (NUM_IMAGES/2 - 2); i >= 0; i--) {
		corrH[i] = H[i] * corrH[i + 1];
	}
	corrH[NUM_IMAGES/2 + 1] = H[NUM_IMAGES/2];
	for(i = (NUM_IMAGES/2 + 2); i < NUM_IMAGES; i++) {
		corrH[i] = corrH[i - 1] * H[i - 1];
	}

	//	corrH[3] = H[3];
	//	corrH[2] = H[2] * corrH[3];
	//	corrH[1] = H[1] * corrH[2];
	//	corrH[0] = H[0] * corrH[1];
	//	corrH[5] = H[4];
	//	corrH[6] = corrH[5] * H[5];
	//	corrH[7] = corrH[6] * H[6];
	//	corrH[8] = corrH[7] * H[7];


	//apply offset to homographys to shift images
	for(i = 0; i < NUM_IMAGES/2; i++) {
		corrH[i] = corrH[i].inv() * offset;
		corrH[NUM_IMAGES - 1 - i] = corrH[NUM_IMAGES - 1 - i] * offset;
	}

	//	corrH[0] = corrH[0].inv() * offset;
	//	corrH[1] = corrH[1].inv() * offset;
	//	corrH[2] = corrH[2].inv() * offset;
	//	corrH[3] = corrH[3].inv() * offset;
	//	corrH[5] = corrH[5] * offset;
	//	corrH[6] = corrH[6] * offset;
	//	corrH[7] = corrH[7] * offset;
	//	corrH[8] = corrH[8] * offset;

	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2);
	printf("Time to manipulate %d homographys with respect to the center image takes ,%.0f, ms\n", NUM_IMAGES - 1, duration);

	blankMask = Mat::ones(Size(images[0].cols, images[0].rows), CV_8U) * 255;
	finalImage = Mat::zeros(Size(IMAGE_WIDTH_RATIO * images[0].cols, images[0].rows), 16);

	gettimeofday(& Time1, NULL);
#pragma omp parallel for
	for (i = 0; i < NUM_IMAGES; i++) {
		//ignore center image
		if (i != NUM_IMAGES / 2) {
			//apply homography to images
			cv::warpPerspective(images[i], correctedImages[i], corrH[i], Size(IMAGE_WIDTH_RATIO * images[0].cols, IMAGE_HEIGHT_RATIO * images[0].rows), WARP_INVERSE_MAP);
			//creates masks for images
			cv::warpPerspective(blankMask, imageMasks[i], corrH[i], Size(IMAGE_WIDTH_RATIO * images[0].cols, IMAGE_HEIGHT_RATIO * images[0].rows), WARP_INVERSE_MAP);
		}
	}
	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2);
	printf("Time to apply homoraphy to %d images takes ,%.0f, ms\n", NUM_IMAGES - 1, duration);

	//copy all projective free images to create a mosaic
	gettimeofday(& Time1, NULL);
	for(i = 0; i < NUM_IMAGES / 2; i++) {
		correctedImages[i].copyTo(finalImage, imageMasks[i]);
		correctedImages[NUM_IMAGES - 1 - i].copyTo(finalImage, imageMasks[NUM_IMAGES - 1 - i]);
	}
	images[NUM_IMAGES / 2].copyTo(finalImage(cv::Rect(cv::Point2f(OFF_WIDTH * images[0].cols, OFF_HEIGHT * images[0].rows),images[0].size())));
	//	correctedImages[0].copyTo(finalImage, imageMasks[0]);
	//	correctedImages[8].copyTo(finalImage, imageMasks[8]);
	//	correctedImages[1].copyTo(finalImage, imageMasks[1]);
	//	correctedImages[7].copyTo(finalImage, imageMasks[7]);
	//	correctedImages[2].copyTo(finalImage, imageMasks[2]);
	//	correctedImages[6].copyTo(finalImage, imageMasks[6]);
	//	correctedImages[3].copyTo(finalImage, imageMasks[3]);
	//	correctedImages[5].copyTo(finalImage, imageMasks[5]);
	//	images[4].copyTo(finalImage(cv::Rect(cv::Point2f(4 * images[0].cols, images[0].rows / 2),images[0].size())));

	myString1 = "matchImagePair";
	myString2 = ".jpg";
	for (i = 0; i < NUM_IMAGES - 1; i++) {
		std::stringstream out;
		out << (i + 1);
		myString3 = myString1 + out.str() + myString2;
		imwrite(myString3, matchImages[i]);
	}
	imwrite("Mosiac.jpg", finalImage);

	gettimeofday(& Time2, NULL);
	duration = timeDiff(Time1, Time2);
	printf("Time to generate images takes ,%.0f, ms\n", duration);
	gettimeofday(& EndTime, NULL);
	duration = timeDiff(StartTime, EndTime);
	printf("Time to run the whole program takes ,%.0f, ms\n", duration);


	//	 namedWindow("finalImage");
	//	 imshow("finalImage", correctedImages[8]);
	//	 char c = ' ';
	//	 while ((c = waitKey(0)) != 'q');
	//Keep window there until user presses 'q' to quit.

	return EXIT_SUCCESS;
}

/*
 * Calculate time duration between two clock time
 */
double timeDiff(struct timeval t1, struct timeval t2)
{
	double time_diff;
	double sec_offset = 0;
	double t1_sec = t1.tv_sec;
	double t2_sec = t2.tv_sec;
	double t1_usec = t1.tv_usec;
	double t2_usec = t2.tv_usec;
	//printf("%f %f %f %f\n", t1_sec, t2_sec, t1_usec, t2_usec);
	if (t2_usec < t1_usec) {
		time_diff = 1e6 + t2_usec - t1_usec;
		sec_offset = -1;
	} else {
		time_diff = t2_usec - t1_usec;
	}

	time_diff = (t2_sec - t1_sec) * 1e6 + sec_offset + time_diff;
	return time_diff / 1e3; //in second
}


/*
 * Look up matching index of interest points from matches, and store
 * the corresponding points into srcPoints and destPoints
 */
void FindMatching(vector<vector<DMatch> > matches, vector<KeyPoint> keypoints1,
		vector<KeyPoint> keypoints2, vector<Point2f>* srcPoints,
		vector<Point2f>* destPoints) {
	int i, j;
	int queryIndex;
	int trainIndex;
	for (i = 0; i < (int) matches.size(); i++) {
		//j is either 0, 1, indicating no match or best match
		for (j = 0; j < (int) matches.at(i).size(); j++) {
			queryIndex = matches.at(i).at(0).queryIdx;
			trainIndex = matches.at(i).at(0).trainIdx;
			srcPoints->push_back(keypoints1.at(queryIndex).pt);
			destPoints->push_back(keypoints2.at(trainIndex).pt);
		}
	}
}
/*
 * Calculate homography matrix
 */
Mat CalculateHomography(int n, vector<Point2f>* p1, vector<Point2f>* p2) {
	int i;
	Mat A = Mat::zeros(2 * n, 9, CV_64FC1);
	Mat U = Mat::zeros(2 * n, 2 * n, CV_64FC1);
	Mat D = Mat::zeros(2 * n, 9, CV_64FC1);
	Mat V = Mat::zeros(9, 9, CV_64FC1);
	Mat H = Mat::zeros(3, 3, CV_64FC1);
	//Set A

	for (i = 0; i < n; i++) {
		A.at<double>(2 * i, 3) = -p1->at(i).x;
		A.at<double>(2 * i, 4) = -p1->at(i).y;
		A.at<double>(2 * i, 5) = -1;
		A.at<double>(2 * i, 6) = p2->at(i).y * p1->at(i).x;
		A.at<double>(2 * i, 7) = p2->at(i).y * p1->at(i).y;
		A.at<double>(2 * i, 8) = p2->at(i).y;
		A.at<double>(2 * i + 1, 0) = p1->at(i).x;
		A.at<double>(2 * i + 1, 1) = p1->at(i).y;
		A.at<double>(2 * i + 1, 2) = 1;
		A.at<double>(2 * i + 1, 6) = -p2->at(i).x * p1->at(i).x;
		A.at<double>(2 * i + 1, 7) = -p2->at(i).x * p1->at(i).y;
		A.at<double>(2 * i + 1, 8) = -p2->at(i).x;
	}

	cv::SVD::compute(A, D, U, V);
	for (i = 0; i < 9; i++) {
		H.at<double>(i / 3, i % 3) = V.at<double>(8, i);
	}
	H = H / H.at<double>(2, 2);
	return H;
}

/*
 *  input argument:
 *  decisionThreshold - maximum pixel distance between the estimate point and the actual point
 *  					to be considered as "inliners"
 *  numberOfRandomSeeds(at least 5) - number of random point pairs used in RANSAC
 *  numberOfTrials - number of iterations ran to decide best correspondence
 * 	acceptThresholdRatio(0-1) - minimum number of "inliners" to be considered as valid correspondence
 *
 */
Mat PerformRANSAC(int ID, double decisionThreshold, int numberOfRandomSeeds,
		int numberOfTrials, float acceptThresholdRatio,
		const vector<Point2f>* srcPoints, const vector<Point2f>* destPoints,
		vector<Point2f>* inlinerSrcPoints, vector<Point2f>* inlinerDestPoints) {
	int i, j, k;
	bool done = false;
	int random[numberOfRandomSeeds];
	int numAccept;
	int sampleSize;
	int acceptThreshold;
	vector<Point2f> src;
	vector<Point2f> dest;
	vector<Point2f> acceptedSrc;
	vector<Point2f> acceptedDest;
	vector<int> acceptedValue;
	Point2f srcP;
	Point2f destP;
	Mat srcP_mat;
	Mat destP_mat;
	double y_x, y_y, distance;
	int bestAcceptSoFar = 0;
	Mat bestTrialHSoFar;
	Mat TrialH;

	sampleSize = srcPoints->size();
	if (sampleSize != (int) destPoints->size()) {
		printf("Error! srcPoints size and desPoints size do not match.\n");
		return Mat::zeros(0, 0, CV_32FC1);
	}
	acceptThreshold = sampleSize * acceptThresholdRatio;

	//Transform source points in to a matrix(3 X sampleSize)
	srcP_mat = Mat::zeros(3, sampleSize, CV_64FC1);
	//destP_mat = Mat::zeros(3, sampleSize, CV_64FC1);
#pragma omp parallel for
	for (i = 0; i < sampleSize; i++) {
		srcP_mat.at<double>(0, i) = srcPoints->at(i).x;
		srcP_mat.at<double>(1, i) = srcPoints->at(i).y;
		srcP_mat.at<double>(2, i) = 1.0;
	}

	//Do N trials
#pragma omp parallel firstprivate(random) private(j, k, done, numAccept, src, dest, TrialH, destP_mat, y_x, y_y, distance)
	{
		src.resize(numberOfRandomSeeds);
		dest.resize(numberOfRandomSeeds);
#pragma omp for
		for (i = 0; i < numberOfTrials; i++) {
			done = false;
			numAccept = 0;
			while (!done) {
				for (j = 0; j < numberOfRandomSeeds; j++) {
					random[j] = rand() % sampleSize;
				}
				done = true;
				for (j = 0; j < numberOfRandomSeeds - 1; j++) {
					for (k = j + 1; k < numberOfRandomSeeds; k++) {
						if (random[j] == random[k]) {
							done = false;
						}
					}
				}
			}
			for (j = 0; j < numberOfRandomSeeds; j++) {
				src.push_back(srcPoints->at(random[j]));
				dest.push_back(destPoints->at(random[j]));
			}

			TrialH = CalculateHomography(src.size(), &src, &dest);
			destP_mat = TrialH * srcP_mat;

			//Compute Estimated Destination points
			for (j = 0; j < sampleSize; j++) {
				y_x = destP_mat.at<double>(0, j) / destP_mat.at<double>(2, j);
				y_y = destP_mat.at<double>(1, j) / destP_mat.at<double>(2, j);
				distance = sqrt(
						pow((y_y - destPoints->at(j).y), 2)
						+ pow((y_x - destPoints->at(j).x), 2));
				if (distance < decisionThreshold) {
					numAccept = numAccept + 1;
				}
			}
#pragma omp critical (selectBestH)
			{
				if (numAccept > acceptThreshold) {
					if (numAccept > bestAcceptSoFar) {
						bestAcceptSoFar = numAccept;
						bestTrialHSoFar = TrialH;
					}
				}
			}
			src.clear();
			dest.clear();
		}
	}
	//Error check
	if (bestAcceptSoFar == 0) {
		printf("Cannot find any accepted Trial\n");
		exit(0);
	}
	printf("%d: acceptance ratio is %d/%d\n", ID, bestAcceptSoFar, sampleSize);
	//Use all inliners to calculate best H
	destP_mat = bestTrialHSoFar * srcP_mat;
#pragma omp parallel for private(srcP, destP, y_x, y_y, distance)
	for (i = 0; i < sampleSize; i++) {
		srcP = srcPoints->at(i);
		destP = destPoints->at(i);
		y_x = destP_mat.at<double>(0, i) / destP_mat.at<double>(2, i);
		y_y = destP_mat.at<double>(1, i) / destP_mat.at<double>(2, i);
		distance = sqrt(pow((y_y - destP.y), 2) + pow((y_x - destP.x), 2));
#pragma omp critical (pickinliners)
		{
			if (distance < decisionThreshold) {
				src.push_back(srcP);
				dest.push_back(destP);
			}
		}
	}

	for (i = 0; i < (int) src.size(); i++) {
		inlinerSrcPoints->push_back(src.at(i));
		inlinerDestPoints->push_back(dest.at(i));
	}
	TrialH = CalculateHomography(src.size(), &src, &dest);
	return TrialH;
}
Mat CreateMatchingImage(Mat image1, Mat image2,
		const vector<Point2f>* srcPoints, const vector<Point2f>* destPoints,
		const vector<Point2f>* inlinerSrcPoints,
		const vector<Point2f>* inlinerDestPoints) {
	int i;
	Point pt;
	Mat matchImage;
	matchImage.create(image1.rows, 2 * image1.cols, image1.type());
	image1.copyTo(matchImage.rowRange(0, image1.rows).colRange(0, image1.cols));
	image2.copyTo(
			matchImage.rowRange(0, image1.rows).colRange(image1.cols,
					2 * image1.cols));
	for (i = 0; i < (int) srcPoints->size(); i++) {
		pt.x = destPoints->at(i).x + image1.cols;
		pt.y = destPoints->at(i).y;
		cv::circle(matchImage, srcPoints->at(i), 3, cvScalar(0, 0, 255), 1);
		cv::circle(matchImage, pt, 3, cvScalar(0, 0, 255), 1);
		cv::line(matchImage, srcPoints->at(i), pt, cvScalar(0, 0, 255));
	}
	for (i = 0; i < (int) inlinerSrcPoints->size(); i++) {
		pt.x = inlinerDestPoints->at(i).x + image1.cols;
		pt.y = inlinerDestPoints->at(i).y;
		cv::circle(matchImage, inlinerSrcPoints->at(i), 3, cvScalar(0, 255, 0),
				1);
		cv::circle(matchImage, pt, 3, cvScalar(0, 255, 0), 1);
		cv::line(matchImage, inlinerSrcPoints->at(i), pt, cvScalar(0, 255, 0));
	}
	return matchImage;
}
